//! 温池管理
//! 
//! 预热但未完全初始化的沙箱，激活时间~50ms

use crate::Result;
use crate::types::{PoolStats, PoolType, Language, SandboxState};
use crate::sandbox_manager::{Sandbox, LifecycleManager};
use crate::pool::cold_pool::ColdPool;
use std::sync::Arc;
use std::collections::HashMap;
use tokio::sync::RwLock;
use std::collections::VecDeque;
use std::time::Instant;
use tracing::{info, warn, error, debug};
use tokio::time::{sleep, Duration};

/// 预热任务信息
#[derive(Debug, Clone)]
pub struct WarmingTask {
    pub language: Language,
    pub target_count: usize,
    pub started_at: std::time::SystemTime,
    pub completed_count: usize,
}

/// 温池管理器
pub struct WarmPool {
    config: Arc<crate::config::PoolConfig>,
    /// 温池存储映射（按语言分组的预热沙箱队列）
    pool: Arc<RwLock<HashMap<Language, VecDeque<Sandbox>>>>,
    /// 冷池引用（用于从冷池创建沙箱到温池）
    cold_pool: Arc<ColdPool>,
    /// 生命周期管理器
    lifecycle_manager: Arc<LifecycleManager>,
    /// 正在预热的任务
    warming_tasks: Arc<RwLock<HashMap<Language, WarmingTask>>>,
    /// 统计信息
    stats: Arc<RwLock<WarmPoolStats>>,
}

/// 温池统计信息
#[derive(Debug, Clone)]
pub struct WarmPoolStats {
    pub total_warmed: usize,
    pub total_activated: usize,
    pub avg_activation_time_ms: f64,
    pub error_count: usize,
}

impl Default for WarmPoolStats {
    fn default() -> Self {
        Self {
            total_warmed: 0,
            total_activated: 0,
            avg_activation_time_ms: 0.0,
            error_count: 0,
        }
    }
}

impl WarmPool {
    /// 创建温池管理器
    pub fn new(
        config: Arc<crate::config::PoolConfig>,
        cold_pool: Arc<ColdPool>,
        lifecycle_manager: Arc<LifecycleManager>,
    ) -> Self {
        Self {
            config,
            pool: Arc::new(RwLock::new(HashMap::new())),
            cold_pool,
            lifecycle_manager,
            warming_tasks: Arc::new(RwLock::new(HashMap::new())),
            stats: Arc::new(RwLock::new(WarmPoolStats::default())),
        }
    }
    
    /// 获取并激活温池沙箱
    pub async fn acquire(&self, language: &Language) -> Result<Option<Sandbox>> {
        let start_time = Instant::now();
        debug!("开始从温池获取沙箱 - 语言: {:?}", language);
        
        let mut sandbox = {
            let mut pool = self.pool.write().await;
            
            if let Some(lang_pool) = pool.get_mut(language) {
                lang_pool.pop_front()
            } else {
                None
            }
        };
        
        if let Some(ref mut sandbox) = sandbox {
            // 激活沙箱（从预热状态转为运行状态）
            let activation_result = self.activate_sandbox(sandbox).await;
            
            match activation_result {
                Ok(()) => {
                    let duration = start_time.elapsed();
                    info!(
                        "从温池激活沙箱成功 - 语言: {:?}, 耗时: {}ms",
                        language,
                        duration.as_millis()
                    );
                    
                    // 更新统计信息
                    let mut stats = self.stats.write().await;
                    stats.total_activated += 1;
                    let new_avg = (stats.avg_activation_time_ms * (stats.total_activated - 1) as f64 
                        + duration.as_millis() as f64) / stats.total_activated as f64;
                    stats.avg_activation_time_ms = new_avg;
                    
                    return Ok(Some(sandbox.clone()));
                }
                Err(e) => {
                    error!("激活温池沙箱失败 - 语言: {:?}, 错误: {}", language, e);
                    let mut stats = self.stats.write().await;
                    stats.error_count += 1;
                    return Err(e);
                }
            }
        }
        
        debug!("温池中没有可用的 {:?} 沙箱", language);
        Ok(None)
    }
    
    /// 向温池添加预热沙箱
    pub async fn add_prewarmed(&self, sandbox: Sandbox) -> Result<()> {
        let language = sandbox.language().clone();
        debug!("向温池添加预热沙箱 - 语言: {:?}", language);
        
        // 验证沙箱状态是否适合放入温池
        if !matches!(sandbox.metadata.state, SandboxState::Stopped | SandboxState::Creating) {
            warn!("沙箱状态不适合放入温池: {:?}", sandbox.metadata.state);
            return Ok(());
        }
        
        let mut pool = self.pool.write().await;
        let lang_pool = pool.entry(language.clone()).or_insert_with(VecDeque::new);
        
        // 检查温池容量限制
        if lang_pool.len() < self.config.warm_pool_size {
            lang_pool.push_back(sandbox);
            debug!("成功添加沙箱到温池 - 语言: {:?}", language);
            
            // 更新统计
            let mut stats = self.stats.write().await;
            stats.total_warmed += 1;
        } else {
            warn!("温池已满，无法添加更多 {:?} 沙箱", language);
        }
        
        Ok(())
    }
    
    /// 预热指定语言的沙箱
    pub async fn preheat(&self, language: Language, target_count: usize) -> Result<()> {
        info!("开始预热温池 - 语言: {:?}, 目标数量: {}", language, target_count);
        
        // 检查是否已经在预热
        {
            let warming_tasks = self.warming_tasks.read().await;
            if warming_tasks.contains_key(&language) {
                warn!("语言 {:?} 已经在预热中", language);
                return Ok(());
            }
        }
        
        // 记录预热任务
        {
            let mut warming_tasks = self.warming_tasks.write().await;
            warming_tasks.insert(language.clone(), WarmingTask {
                language: language.clone(),
                target_count,
                started_at: std::time::SystemTime::now(),
                completed_count: 0,
            });
        }
        
        // 获取当前温池中该语言的沙箱数量
        let current_count = {
            let pool = self.pool.read().await;
            pool.get(&language).map(|q| q.len()).unwrap_or(0)
        };
        
        let needed_count = target_count.saturating_sub(current_count);
        if needed_count == 0 {
            info!("温池已有足够的 {:?} 沙箱，无需预热", language);
            self.warming_tasks.write().await.remove(&language);
            return Ok(());
        }
        
        info!("需要预热 {} 个 {:?} 沙箱", needed_count, language);
        
        // 启动预热任务
        let cold_pool = Arc::clone(&self.cold_pool);
        let lifecycle_manager = Arc::clone(&self.lifecycle_manager);
        let warm_pool = Arc::new(self.clone_for_task());
        let warming_tasks = Arc::clone(&self.warming_tasks);
        
        tokio::spawn(async move {
            Self::execute_warming_task(
                cold_pool,
                lifecycle_manager,
                warm_pool,
                warming_tasks,
                language,
                needed_count,
            ).await;
        });
        
        Ok(())
    }
    
    /// 获取温池统计
    pub async fn get_stats(&self) -> Result<PoolStats> {
        let pool = self.pool.read().await;
        let idle_count: usize = pool.values().map(|v| v.len()).sum();
        let warming_count = self.warming_tasks.read().await.len();
        let stats = self.stats.read().await;
        
        Ok(PoolStats {
            pool_type: PoolType::Warm,
            total_capacity: self.config.warm_pool_size * pool.len(),
            used_count: 0, // 温池中的都是预热状态，不算使用中
            idle_count,
            creating_count: warming_count,
            error_count: stats.error_count,
            avg_allocation_latency_ms: stats.avg_activation_time_ms,
            cache_hit_rate: if idle_count > 0 { 90.0 } else { 0.0 },
        })
    }
    
    /// 获取特定语言的沙箱数量
    pub async fn sandbox_count(&self, language: &Language) -> usize {
        let pool = self.pool.read().await;
        pool.get(language).map(|q| q.len()).unwrap_or(0)
    }
    
    /// 清理超时的预热沙箱
    pub async fn cleanup_expired_sandboxes(&self) -> Result<()> {
        let mut pool = self.pool.write().await;
        let max_idle_time = Duration::from_secs(self.config.max_idle_seconds);
        let now = std::time::SystemTime::now();
        
        for (language, sandbox_queue) in pool.iter_mut() {
            let initial_count = sandbox_queue.len();
            
            // 保留未超时的沙箱
            let mut new_queue = VecDeque::new();
            while let Some(sandbox) = sandbox_queue.pop_front() {
                let idle_duration = sandbox.metadata.idle_time()
                    .unwrap_or(Duration::from_secs(0));
                    
                if idle_duration < max_idle_time {
                    new_queue.push_back(sandbox);
                } else {
                    debug!("清理超时的温池沙箱 - 语言: {:?}", language);
                }
            }
            
            *sandbox_queue = new_queue;
            
            let removed_count = initial_count - sandbox_queue.len();
            if removed_count > 0 {
                info!("清理超时温池沙箱 - 语言: {:?}, 清理数量: {}", language, removed_count);
            }
        }
        
        Ok(())
    }
    
    /// 启动自动预热任务
    pub async fn start_auto_preheating(&self) -> Result<()> {
        info!("启动自动预热任务");
        
        let warm_pool = Arc::new(self.clone_for_task());
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(30));
            
            loop {
                interval.tick().await;
                
                // 检查各语言的温池状态并决定是否需要预热
                if let Err(e) = warm_pool.check_and_preheat().await {
                    error!("自动预热检查失败: {}", e);
                }
            }
        });
        
        Ok(())
    }
    
    // 私有方法
    
    /// 激活沙箱（从预热状态转为运行状态）
    async fn activate_sandbox(&self, sandbox: &mut Sandbox) -> Result<()> {
        debug!("激活预热沙箱: {}", sandbox.metadata.id);
        
        // 启动沙箱
        self.lifecycle_manager.start_sandbox(&mut sandbox.metadata).await?;
        
        // 等待短暂时间确保完全激活
        sleep(Duration::from_millis(10)).await;
        
        debug!("沙箱激活完成: {}", sandbox.metadata.id);
        Ok(())
    }
    
    /// 执行预热任务
    async fn execute_warming_task(
        cold_pool: Arc<ColdPool>,
        lifecycle_manager: Arc<LifecycleManager>,
        warm_pool: Arc<WarmPool>,
        warming_tasks: Arc<RwLock<HashMap<Language, WarmingTask>>>,
        language: Language,
        count: usize,
    ) {
        info!("执行预热任务 - 语言: {:?}, 数量: {}", language, count);
        
        let mut successful_count = 0;
        
        for i in 0..count {
            // 从冷池创建沙箱
            match cold_pool.acquire(&language).await {
                Ok(Some(sandbox)) => {
                    // 将沙箱添加到温池
                    if let Err(e) = warm_pool.add_prewarmed(sandbox).await {
                        error!("添加预热沙箱到温池失败: {}", e);
                    } else {
                        successful_count += 1;
                        
                        // 更新任务进度
                        {
                            let mut tasks = warming_tasks.write().await;
                            if let Some(task) = tasks.get_mut(&language) {
                                task.completed_count = successful_count;
                            }
                        }
                        
                        debug!("预热进度: {}/{} - 语言: {:?}", successful_count, count, language);
                    }
                }
                Ok(None) => {
                    warn!("冷池无法提供 {:?} 沙箱，预热任务 {}/{} 跳过", language, i + 1, count);
                }
                Err(e) => {
                    error!("从冷池创建沙箱失败: {}", e);
                }
            }
            
            // 控制预热速度，避免资源争抢
            sleep(Duration::from_millis(100)).await;
        }
        
        // 移除完成的预热任务
        warming_tasks.write().await.remove(&language);
        
        info!(
            "预热任务完成 - 语言: {:?}, 成功: {}/{}", 
            language, successful_count, count
        );
    }
    
    /// 检查并执行自动预热
    async fn check_and_preheat(&self) -> Result<()> {
        let pool = self.pool.read().await;
        
        for language in [Language::Python, Language::JavaScript, Language::Rust, Language::Go, Language::Java] {
            let current_count = pool.get(&language).map(|q| q.len()).unwrap_or(0);
            let target_count = self.config.warm_pool_size / 2; // 维持一半容量
            
            if current_count < target_count {
                drop(pool); // 释放读锁
                
                let preheat_count = target_count - current_count;
                info!(
                    "自动预热触发 - 语言: {:?}, 当前: {}, 目标: {}, 需要预热: {}", 
                    language, current_count, target_count, preheat_count
                );
                
                if let Err(e) = self.preheat(language.clone(), target_count).await {
                    error!("自动预热失败 - 语言: {:?}, 错误: {}", language, e);
                }
                
                return Ok(()); // 一次只处理一种语言，避免资源争抢
            }
        }
        
        Ok(())
    }
    
    /// 为任务克隆温池（共享数据）
    fn clone_for_task(&self) -> WarmPool {
        WarmPool {
            config: Arc::clone(&self.config),
            pool: Arc::clone(&self.pool),
            cold_pool: Arc::clone(&self.cold_pool),
            lifecycle_manager: Arc::clone(&self.lifecycle_manager),
            warming_tasks: Arc::clone(&self.warming_tasks),
            stats: Arc::clone(&self.stats),
        }
    }
}